#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Chim -> Zig代码生成器
生成可编译的Zig代码
"""

from parser_ import *
from typing import List

# 导入组生命周期AST节点
try:
    from parser_ import (GroupBlock, SnapshotExpression, HandleExpression, ArenaCall,
                         PointerType, AddressOf, Dereference, InlineAssembly)
except ImportError:
    # 兼容旧版本
    GroupBlock = None
    SnapshotExpression = None
    HandleExpression = None
    ArenaCall = None
    PointerType = None
    AddressOf = None
    Dereference = None
    InlineAssembly = None

class ZigCodeGenerator:
    """Zig代码生成器"""
    
    def __init__(self, dialect: str = "safe"):
        self.dialect = dialect
        self.code_lines: List[str] = []
        self.indent_level = 0
        self.indent_str = "    "  # Zig使用4空格缩进
    
    def emit(self, line: str = ""):
        """生成一行代码"""
        if line:
            self.code_lines.append(self.indent_str * self.indent_level + line)
        else:
            self.code_lines.append("")
    
    def generate(self, ast) -> str:
        """生成Zig代码"""
        self.code_lines.clear()
        
        # 处理Program对象
        if hasattr(ast, 'statements'):
            statements = ast.statements
        elif isinstance(ast, list):
            statements = ast
        else:
            statements = [ast]
        
        # 头部
        self.emit("// Generated by Chim Compiler")
        self.emit("// Zig Backend")
        self.emit()
        self.emit("const std = @import(\"std\");")
        self.emit()
        
        # 生成所有语句
        for stmt in statements:
            self.visit_statement(stmt)
            self.emit()
        
        # 添加main函数（如果没有）
        has_main = any(isinstance(s, Function) and s.name in ("主", "main", "主函数") 
                       for s in statements)
        
        if not has_main:
            self.emit("pub fn main() !void {")
            self.indent_level += 1
            self.emit("// No main function defined")
            self.indent_level -= 1
            self.emit("}")
        
        return "\n".join(self.code_lines)
    
    def visit_statement(self, stmt: ASTNode):
        """访问语句"""
        if isinstance(stmt, Function):
            self.visit_function(stmt)
        elif isinstance(stmt, Struct):
            self.visit_struct(stmt)
        elif isinstance(stmt, VariableDeclaration):
            self.visit_variable_declaration(stmt)
        elif isinstance(stmt, MatchStatement):
            self.visit_match_statement(stmt)
        elif isinstance(stmt, ForLoop):
            self.visit_for_loop(stmt)
        elif isinstance(stmt, ReturnStatement):
            self.visit_return_statement(stmt)
        elif isinstance(stmt, CallExpression):
            self.emit(self.visit_expression(stmt) + ";")
        elif isinstance(stmt, InlineAssembly):
            self.visit_inline_assembly(stmt)
        elif GroupBlock and isinstance(stmt, GroupBlock):
            self.visit_group_block(stmt)
    
    def visit_function(self, node: Function):
        """访问函数定义"""
        # 函数名映射
        func_name = self.map_function_name(node.name)
        
        # 参数列表
        params = []
        for param_name, param_type in node.params:
            zig_type = self.map_type(param_type)
            params.append(f"{param_name}: {zig_type}")
        
        params_str = ", ".join(params)
        
        # 返回类型
        return_type = self.map_type(node.return_type) if node.return_type else "void"
        
        # pub for main
        pub = "pub " if func_name == "main" else ""
        
        self.emit(f"{pub}fn {func_name}({params_str}) {return_type} {{")
        self.indent_level += 1
        
        # 函数体
        for stmt in node.body:
            self.visit_statement(stmt)
        
        # 如果没有显式return且返回void，无需添加
        
        self.indent_level -= 1
        self.emit("}")
    
    def visit_struct(self, node: Struct):
        """访问结构体定义"""
        self.emit(f"const {node.name} = struct {{")
        self.indent_level += 1
        
        for field_name, field_type in node.fields:
            zig_type = self.map_type(field_type)
            self.emit(f"{field_name}: {zig_type},")
        
        self.indent_level -= 1
        self.emit("};")
    
    def visit_variable_declaration(self, node: VariableDeclaration):
        """访问变量声明"""
        keyword = "var" if node.is_mutable else "const"
        value_str = self.visit_expression(node.value)
        
        if node.type:
            zig_type = self.map_type(node.type)
            self.emit(f"{keyword} {node.name}: {zig_type} = {value_str};")
        else:
            self.emit(f"{keyword} {node.name} = {value_str};")
    
    def visit_match_statement(self, node: MatchStatement):
        """访问match语句 - 转换为switch或if-else"""
        expr_str = self.visit_expression(node.expression)
        
        # Zig的switch要求编译时已知所有case
        # 对于整数等可以用switch，否则用if-else链
        
        # 简化: 使用if-else链（更通用）
        # 使用临时变量保存表达式值
        temp_var = "_match_value"
        self.emit(f"const {temp_var} = {expr_str};")
        
        first = True
        
        for item in node.cases:
            if len(item) == 3:
                pattern, guard, body = item
            else:
                pattern, body = item
                guard = None
            
            # 检查默认分支
            is_default = isinstance(pattern, Literal) and pattern.value_type == "DEFAULT"
            is_underscore = isinstance(pattern, Identifier) and pattern.name == "_"
            
            if is_default or is_underscore:
                self.emit("} else {")
            else:
                # 绑定模式变量
                bindings = self.generate_pattern_bindings(temp_var, pattern)
                
                # 生成条件
                condition = self.generate_match_condition(temp_var, pattern)
                
                if guard:
                    guard_str = self.visit_expression(guard)
                    condition = f"({condition}) and ({guard_str})"
                
                if first:
                    self.emit(f"if ({condition}) {{")
                    first = False
                else:
                    self.emit(f"}} else if ({condition}) {{")
                
                # 在case块内生成绑定
                self.indent_level += 1
                for binding in bindings:
                    self.emit(binding)
            
            # case体
            if not (is_default or is_underscore):
                # 已经增加了缩进
                pass
            else:
                self.indent_level += 1
            
            for stmt in body:
                self.visit_statement(stmt)
            self.indent_level -= 1
        
        self.emit("}")
    
    def generate_pattern_bindings(self, temp_var: str, pattern: ASTNode) -> List[str]:
        """生成模式变量绑定 - Zig"""
        bindings = []
        
        if isinstance(pattern, TuplePattern):
            for i, elem in enumerate(pattern.elements):
                if isinstance(elem, Identifier):
                    var_name = elem.name
                    bindings.append(f"const {var_name} = {temp_var}[{i}];")
        elif isinstance(pattern, Identifier) and pattern.name not in ('_', '默认', 'true', 'false', '真', '假'):
            # 单个模式变量
            var_name = pattern.name
            bindings.append(f"const {var_name} = {temp_var};")
        
        return bindings
    
    def generate_match_condition(self, expr_str: str, pattern: ASTNode) -> str:
        """生成匹配条件"""
        if isinstance(pattern, TuplePattern):
            # 元组匹配 - Zig需要特殊处理
            conditions = []
            for i, elem in enumerate(pattern.elements):
                if isinstance(elem, Identifier):
                    # 模式变量 - 任意匹配
                    continue
                else:
                    value_str = self.visit_expression(elem)
                    conditions.append(f"{expr_str}[{i}] == {value_str}")
            return " and ".join(conditions) if conditions else "true"
        
        elif isinstance(pattern, RangePattern):
            start_str = self.visit_expression(pattern.start)
            end_str = self.visit_expression(pattern.end)
            if pattern.inclusive:
                return f"({expr_str} >= {start_str} and {expr_str} <= {end_str})"
            else:
                return f"({expr_str} >= {start_str} and {expr_str} < {end_str})"
        
        else:
            pattern_str = self.visit_expression(pattern)
            return f"{expr_str} == {pattern_str}"
    
    def visit_for_loop(self, node: ForLoop):
        """访问for循环"""
        iterable_str = self.visit_expression(node.iterable)
        
        # Zig的for循环
        self.emit(f"for ({iterable_str}) |{node.variable}| {{")
        self.indent_level += 1
        
        for stmt in node.body:
            self.visit_statement(stmt)
        
        self.indent_level -= 1
        self.emit("}")
    
    def visit_return_statement(self, node: ReturnStatement):
        """访闞return语句"""
        if node.expression:
            expr_str = self.visit_expression(node.expression)
            self.emit(f"return {expr_str};")
        else:
            self.emit("return;")
        
    def visit_group_block(self, node):
        """访问组块 - 生成Zig的defer机制"""
        group_name = node.name.name if isinstance(node.name, Identifier) else str(node.name)
            
        # Zig使用defer进行清理
        self.emit(f"// 组: {group_name}")
        self.emit("{")
        self.indent_level += 1
            
        # 生成组内语句
        for stmt in node.body:
            self.visit_statement(stmt)
            
        # 添加defer清理
        self.emit(f"// defer: 组 '{group_name}' 结束")
            
        self.indent_level -= 1
        self.emit("}")
    
    def visit_expression(self, expr: ASTNode) -> str:
        """访问表达式"""
        if isinstance(expr, Literal):
            return self.visit_literal(expr)
        elif isinstance(expr, Identifier):
            return expr.name
        elif isinstance(expr, BinaryExpression):
            left = self.visit_expression(expr.left)
            right = self.visit_expression(expr.right)
            op = self.map_operator(expr.operator)
            return f"({left} {op} {right})"
        elif isinstance(expr, UnaryExpression):
            operand = self.visit_expression(expr.operand)
            op = self.map_operator(expr.operator)
            return f"{op}{operand}"
        elif isinstance(expr, CallExpression):
            return self.visit_call_expression(expr)
        elif isinstance(expr, MemberAccess):
            obj = self.visit_expression(expr.object)
            return f"{obj}.{expr.member.name}"
        elif isinstance(expr, ArrayLiteral):
            return self.visit_array_literal(expr)
        elif isinstance(expr, MapLiteral):
            return self.visit_map_literal(expr)
        elif SnapshotExpression and isinstance(expr, SnapshotExpression):
            # 快照 - 生成副本
            target = self.visit_expression(expr.target)
            return f"@as(@TypeOf({target}), {target})  // 快照"
        elif HandleExpression and isinstance(expr, HandleExpression):
            # 句柄 - 直接引用
            target = self.visit_expression(expr.target)
            return f"&{target}  // 句柄"
        elif ArenaCall and isinstance(expr, ArenaCall):
            return "std.heap.page_allocator  // Arena"
        elif AddressOf and isinstance(expr, AddressOf):
            # 取地址: &variable
            target = self.visit_expression(expr.target)
            return f"&{target}"
        elif Dereference and isinstance(expr, Dereference):
            # 解引用: variable.* 或 *variable
            target = self.visit_expression(expr.target)
            return f"{target}.*"
        else:
            return "undefined"
    
    def visit_literal(self, node: Literal) -> str:
        """访问字面量"""
        if node.value_type == "INTEGER":
            return str(node.value)
        elif node.value_type == "FLOAT":
            return str(node.value)
        elif node.value_type == "STRING":
            # Zig字符串需要转义
            return f'"{node.value}"'
        elif node.value_type == "BOOLEAN":
            return "true" if node.value else "false"
        elif node.value_type == "NIL":
            return "null"
        else:
            return "undefined"
    
    def visit_array_literal(self, node: ArrayLiteral) -> str:
        """访问数组字面量"""
        if not node.elements:
            return "[_]i64{}"
        elements = [self.visit_expression(elem) for elem in node.elements]
        return f"[_]i64{{ {', '.join(elements)} }}"
    
    def visit_map_literal(self, node: MapLiteral) -> str:
        """访问映射字面量"""
        # Zig中使用std.StringHashMap或std.AutoHashMap
        return "std.AutoHashMap(i64, i64).init(std.heap.page_allocator)  // Map"
    
    def visit_inline_assembly(self, node):
        """访问内联汇编"""
        vol = "volatile " if node.is_volatile else ""
        
        # Zig 的内联汇编语法
        # asm volatile (code : outputs : inputs : clobbers)
        self.emit(f"_ = asm {vol}(")
        self.indent_level += 1
        
        # 汇编代码 (使用 \\\ 转义字符串)
        code_lines = node.code.split('\\n')
        for line in code_lines:
            self.emit(f"\\\\ {line}")
        
        # TODO: 添加约束支持 (当前简化版本)
        
        self.indent_level -= 1
        self.emit(");")
        self.emit()
    
    def visit_call_expression(self, node: CallExpression) -> str:
        """访问函数调用"""
        func_name = self.visit_expression(node.function)
        args = [self.visit_expression(arg) for arg in node.arguments]
        
        # 特殊处理内置函数
        if func_name in ("输出", "print"):
            # Zig的打印 - 处理多个参数
            if len(args) == 0:
                return 'std.debug.print("\\n", .{})'
            elif len(args) == 1:
                arg = args[0]
                if arg.startswith('"') and arg.endswith('"'):
                    return f'std.debug.print({arg} ++ "\\n", .{{}})'
                else:
                    return f'std.debug.print("{{}}\\n", .{{{arg}}})'
            else:
                # 多个参数
                format_parts = []
                for arg in args:
                    if arg.startswith('"') and arg.endswith('"'):
                        format_parts.append('std.debug.print(' + arg + ' ++ " ", .{})')
                    else:
                        format_parts.append(f'std.debug.print("{{}} ", .{{{arg}}})')
                format_parts.append('std.debug.print("\\n", .{})')
                return '; '.join(format_parts)
        
        # 端口 I/O 内置函数
        if func_name in ("inb", "outb", "inw", "outw", "ind", "outd"):
            return self.generate_port_io(func_name, args)
        
        return f"{func_name}({', '.join(args)})"
    
    def generate_port_io(self, func_name: str, args: list) -> str:
        """生成端口 I/O 代码"""
        if func_name == "inb":
            # inb(port) -> u8
            port = args[0] if args else "0"
            return f"asm volatile (\"inb %[port], %[result]\" : [result] \"={{al}}\" (-> u8) : [port] \"{{dx}}\" ({port}))"
        elif func_name == "outb":
            # outb(port, value)
            port = args[0] if len(args) > 0 else "0"
            value = args[1] if len(args) > 1 else "0"
            return f"asm volatile (\"outb %[value], %[port]\" : : [port] \"{{dx}}\" ({port}), [value] \"{{al}}\" ({value}))"
        elif func_name == "inw":
            # inw(port) -> u16
            port = args[0] if args else "0"
            return f"asm volatile (\"inw %[port], %[result]\" : [result] \"={{ax}}\" (-> u16) : [port] \"{{dx}}\" ({port}))"
        elif func_name == "outw":
            # outw(port, value)
            port = args[0] if len(args) > 0 else "0"
            value = args[1] if len(args) > 1 else "0"
            return f"asm volatile (\"outw %[value], %[port]\" : : [port] \"{{dx}}\" ({port}), [value] \"{{ax}}\" ({value}))"
        elif func_name == "ind":
            # ind(port) -> u32
            port = args[0] if args else "0"
            return f"asm volatile (\"inl %[port], %[result]\" : [result] \"={{eax}}\" (-> u32) : [port] \"{{dx}}\" ({port}))"
        elif func_name == "outd":
            # outd(port, value)
            port = args[0] if len(args) > 0 else "0"
            value = args[1] if len(args) > 1 else "0"
            return f"asm volatile (\"outl %[value], %[port]\" : : [port] \"{{dx}}\" ({port}), [value] \"{{eax}}\" ({value}))"
        
        return f"{func_name}({', '.join(args)})"

    def map_function_name(self, name: str) -> str:
        """映射函数名"""
        mapping = {
            "主": "main",
            "主函数": "main",
        }
        return mapping.get(name, name)
    
    def map_type(self, type_str: str) -> str:
        """映射类型"""
        if not type_str:
            return "void"
        
        mapping = {
            "整数": "i64",
            "浮点数": "f64",
            "字符": "u8",
            "字符串": "[]const u8",
            "布尔": "bool",
            "单元": "void",
            # 英文别名
            "int": "i64",
            "float": "f64",
            "char": "u8",
            "string": "[]const u8",
            "bool": "bool",
            "unit": "void",
        }
        
        return mapping.get(type_str, type_str)
    
    def map_operator(self, op: str) -> str:
        """映射运算符"""
        mapping = {
            "&&": "and",
            "||": "or",
            "!": "!",
            "==": "==",
            "!=": "!=",
        }
        return mapping.get(op, op)

def generate_zig_code(ast, dialect: str = "safe") -> str:
    """生成Zig代码（入口函数）"""
    generator = ZigCodeGenerator(dialect=dialect)
    return generator.generate(ast)
